"""
Battery data processing module for SOC estimation
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from typing import Tuple, Dict, List, Optional
from sklearn.preprocessing import StandardScaler
import os
import scipy.io as sio


def prepare_sequence_data(X: np.ndarray, y: np.ndarray, window_size: int = 10) -> Tuple[np.ndarray, np.ndarray]:
    """
    Prepare data with sliding window for sequence-based SOC estimation
    
    Args:
        X: Feature matrix
        y: Target vector (SOC)
        window_size: Size of the sliding window
        
    Returns:
        Tuple of (windowed_features, target_soc)
    """
    n_samples, n_features = X.shape
    n_windows = n_samples - window_size + 1
    
    # Initialize arrays
    X_windowed = np.zeros((n_windows, window_size * n_features))
    y_windowed = np.zeros(n_windows)
    
    # Create windowed data
    for i in range(n_windows):
        # Extract window
        window = X[i:i + window_size]
        
        # Flatten window
        X_windowed[i] = window.flatten()
        
        # Target is the SOC at the end of the window
        y_windowed[i] = y[i + window_size - 1]
    
    return X_windowed, y_windowed


def load_battery_data(data_dir: str = None, data_file: str = None, data_format: str = 'mat', soc_method: str = 'voltage_based') -> Tuple[np.ndarray, np.ndarray]:
    """
    Load battery data for SOC estimation
    
    Args:
        data_dir: Directory containing data files
        data_file: Specific data file name
        data_format: Data format ('csv' or 'mat')
        soc_method: Method to calculate SOC ('ah_counting', 'voltage_based', 'combined')
        
    Returns:
        Tuple of (features, target)
    """
    # 如果未提供数据目录，使用合成数据
    if data_dir is None:
        print("No data directory provided. Generating synthetic data...")
        return generate_synthetic_data(n_samples=1000)
    
    # 构建文件路径
    if data_file is None:
        # 如果未提供特定文件，尝试查找第一个匹配的文件
        if data_format == 'csv':
            files = [f for f in os.listdir(data_dir) if f.endswith('.csv')]
        else:
            files = [f for f in os.listdir(data_dir) if f.endswith('.mat')]
        
        if not files:
            raise FileNotFoundError(f"No {data_format} files found in {data_dir}")
        
        data_file = files[0]
    
    file_path = os.path.join(data_dir, data_file)
    print(f"Loading data from {file_path}")
    
    # 根据数据格式加载数据
    if data_format == 'csv':
        return load_csv_data(file_path)
    else:
        return load_panasonic_data(file_path, soc_method)


def load_csv_data(file_path: str) -> Tuple[np.ndarray, np.ndarray]:
    """
    Load battery data from CSV file
    
    Args:
        file_path: Path to the CSV file
        
    Returns:
        Tuple of (features, target)
    """
    try:
        # 尝试加载数据
        data = pd.read_csv(file_path)
        print(f"Data loaded successfully. Shape: {data.shape}")
        
        # 假设CSV文件包含特征列和一个名为'SOC'的目标列
        if 'SOC' in data.columns:
            y = data['SOC'].values
            X = data.drop('SOC', axis=1).values
        else:
            # 如果没有SOC列，假设最后一列是目标
            y = data.iloc[:, -1].values
            X = data.iloc[:, :-1].values
        
        return X, y
    except Exception as e:
        print(f"Error loading CSV data: {e}")
        raise


def load_panasonic_data(file_path: str, soc_calculation_method: str = 'ah_counting') -> Tuple[np.ndarray, np.ndarray]:
    """
    Load Panasonic 18650PF battery data from .mat files
    
    Args:
        file_path: Path to the .mat data file
        soc_calculation_method: Method to calculate SOC ('ah_counting', 'voltage_based', 'combined' or 'provided')
        
    Returns:
        Tuple of (features, target)
    """
    print(f"Loading Panasonic 18650PF battery data from {file_path}...")
    
    try:
        # Check file extension
        if not file_path.endswith('.mat'):
            raise ValueError("File must be a .mat file")
        
        # Load .mat file
        mat_data = sio.loadmat(file_path)
        
        # 打印.mat文件中的所有变量名，帮助调试
        print(f"Variables in .mat file: {list(mat_data.keys())}")
        
        # 检查数据结构，处理不同格式的.mat文件
        if 'meas' in mat_data:
            # 处理包含'meas'结构体的.mat文件
            print("Found 'meas' structure in .mat file")
            meas = mat_data['meas']
            
            # 从结构体中提取数据
            # 注意：MATLAB结构体在Python中被加载为包含字段名和值的特殊数组
            # 获取结构体字段名
            field_names = meas.dtype.names
            print(f"Fields in 'meas' structure: {field_names}")
            
            # 详细打印每个字段的前几个值和统计信息
            for field in field_names:
                field_data = _extract_from_struct(meas, field)
                if field_data is not None:
                    # 打印前5个值
                    print(f"Field '{field}' - First 5 values: {field_data[:5]}")
                    
                    # 检查数据类型，只对数值类型计算统计信息
                    if np.issubdtype(field_data.dtype, np.number):
                        # 打印统计信息
                        print(f"Field '{field}' - Shape: {field_data.shape}, "
                              f"Min: {np.nanmin(field_data) if not np.all(np.isnan(field_data)) else 'All NaN'}, "
                              f"Max: {np.nanmax(field_data) if not np.all(np.isnan(field_data)) else 'All NaN'}, "
                              f"NaN count: {np.isnan(field_data).sum()}, "
                              f"Data type: {field_data.dtype}")
                    else:
                        # 对于非数值类型，只打印形状和类型
                        print(f"Field '{field}' - Shape: {field_data.shape}, "
                              f"Data type: {field_data.dtype} (non-numeric)")
                else:
                    print(f"Field '{field}' could not be extracted")
            
            # 提取所需字段
            voltage = _extract_from_struct(meas, 'Voltage')
            current = _extract_from_struct(meas, 'Current')
            ah = _extract_from_struct(meas, 'Ah')
            power = _extract_from_struct(meas, 'Power')
            battery_temp = _extract_from_struct(meas, 'Battery_Temp_degC')
            chamber_temp = _extract_from_struct(meas, 'Chamber_Temp_degC')
            
            # 如果某些字段未找到，尝试其他可能的名称
            if voltage is None:
                voltage = _extract_from_struct(meas, 'voltage')
            if current is None:
                current = _extract_from_struct(meas, 'current')
            if ah is None:
                ah = _extract_from_struct(meas, 'ah')
            if power is None:
                power = _extract_from_struct(meas, 'power')
            if battery_temp is None:
                battery_temp = _extract_from_struct(meas, 'temp')
            if chamber_temp is None:
                chamber_temp = _extract_from_struct(meas, 'ambient')
        else:
            # 尝试直接从.mat文件中提取数组
            voltage = _extract_array(mat_data, ['voltage', 'Voltage', 'V'])
            current = _extract_array(mat_data, ['current', 'Current', 'I'])
            ah = _extract_array(mat_data, ['ah', 'Ah', 'ampere_hour'])
            power = _extract_array(mat_data, ['power', 'Power', 'P'])
            battery_temp = _extract_array(mat_data, ['temp', 'Temp', 'temperature', 'battery_temp'])
            chamber_temp = _extract_array(mat_data, ['chamber_temp', 'ambient_temp', 'ambient'])
        
        # 检查是否成功提取了所有必需的数据
        if voltage is None or current is None or ah is None:
            raise ValueError("Failed to extract required data (voltage, current, or ah) from .mat file")
        
        # 如果power未找到，则从电压和电流计算
        if power is None and voltage is not None and current is not None:
            power = voltage * current
        
        # 如果chamber_temp未找到，使用battery_temp代替
        if chamber_temp is None and battery_temp is not None:
            chamber_temp = battery_temp.copy()
        
        # 打印提取的数据形状
        print(f"Extracted data shapes:")
        print(f"Voltage: {voltage.shape}")
        print(f"Current: {current.shape}")
        print(f"Ah: {ah.shape}")
        print(f"Power: {power.shape if power is not None else 'Not found'}")
        print(f"Battery_Temp: {battery_temp.shape if battery_temp is not None else 'Not found'}")
        print(f"Chamber_Temp: {chamber_temp.shape if chamber_temp is not None else 'Not found'}")
        
        # 检查提取的数据是否包含NaN值
        print(f"NaN values in extracted data:")
        print(f"Voltage NaN count: {np.isnan(voltage).sum()}")
        print(f"Current NaN count: {np.isnan(current).sum()}")
        print(f"Ah NaN count: {np.isnan(ah).sum()}")
        if power is not None:
            print(f"Power NaN count: {np.isnan(power).sum()}")
        if battery_temp is not None:
            print(f"Battery_Temp NaN count: {np.isnan(battery_temp).sum()}")
        if chamber_temp is not None:
            print(f"Chamber_Temp NaN count: {np.isnan(chamber_temp).sum()}")
        
        # Calculate SOC based on selected method
        if soc_calculation_method == 'ah_counting':
            # 改进的安时计数法，使用绝对Ah值而不是相对变化
            # 基于松下18650PF电池的额定容量2.9Ah
            full_capacity = 2.9  # Ah
            
            # 找到最小Ah值（最大放电量）
            min_ah_idx = np.argmin(ah)
            min_ah = ah[min_ah_idx]
            
            # 计算SOC：使用绝对Ah值相对于总容量的比例
            # 注意：这里假设Ah是负值，表示放电量
            soc = 100 * (1 + ah / full_capacity)
            
            print(f"使用改进的安时计数法计算SOC")
            print(f"最小Ah值: {min_ah:.4f}，对应索引: {min_ah_idx}")
            print(f"电池容量: {full_capacity} Ah")
            
        elif soc_calculation_method == 'voltage_based':
            # 基于电压的SOC估计
            soc = np.zeros_like(voltage)
            for i, v in enumerate(voltage):
                soc[i] = _estimate_soc_from_voltage(v)
            
            print(f"使用基于电压的方法估计SOC")
            
        elif soc_calculation_method == 'combined':
            # 结合电压和安时计数的方法
            # 首先使用电压估计初始SOC
            initial_soc = _estimate_soc_from_voltage(voltage[0])
            
            # 然后使用安时计数法计算后续SOC变化
            soc = np.zeros_like(voltage)
            soc[0] = initial_soc
            
            # 使用安时计数法计算SOC变化
            full_capacity = 2.9  # Ah
            for i in range(1, len(voltage)):
                delta_ah = ah[i] - ah[i-1]
                delta_soc = -delta_ah / full_capacity * 100  # 负号是因为放电时Ah增加
                soc[i] = soc[i-1] + delta_soc
            
            # 限制SOC范围在0-100%
            soc = np.clip(soc, 0, 100)
            
            print(f"使用电压和安时计数结合的方法估计SOC")
            print(f"初始SOC（基于电压）: {initial_soc:.2f}%")
            
        else:
            # 如果SOC在数据中提供（不太可能）
            soc = _extract_array(mat_data, ['soc', 'SOC', 'state_of_charge'])
            if soc is None:
                # 如果找不到SOC数据，默认使用电压估计
                print(f"未找到SOC数据，使用电压估计SOC")
                soc = np.zeros_like(voltage)
                for i, v in enumerate(voltage):
                    soc[i] = _estimate_soc_from_voltage(v)
        
        # 打印SOC的统计信息
        print(f"SOC statistics:")
        print(f"SOC shape: {soc.shape}")
        print(f"SOC range: {np.min(soc):.2f}% to {np.max(soc):.2f}%")
        print(f"SOC mean: {np.mean(soc):.2f}%")
        print(f"SOC NaN count: {np.isnan(soc).sum()}")
        
        # 组合特征
        features = []
        feature_names = []
        
        # 添加电压
        if voltage is not None and not np.all(np.isnan(voltage)):
            features.append(voltage)
            feature_names.append('Voltage')
        else:
            print("警告: 电压数据全部为NaN或不存在，将被排除")
        
        # 添加电流
        if current is not None and not np.all(np.isnan(current)):
            features.append(current)
            feature_names.append('Current')
        else:
            print("警告: 电流数据全部为NaN或不存在，将被排除")
        
        # 添加功率（如果有）
        if power is not None and not np.all(np.isnan(power)):
            features.append(power)
            feature_names.append('Power')
        else:
            print("警告: 功率数据全部为NaN或不存在，将被排除")
        
        # 添加电池温度（如果有）
        if battery_temp is not None and not np.all(np.isnan(battery_temp)):
            features.append(battery_temp)
            feature_names.append('Battery_Temp')
        else:
            print("警告: 电池温度数据全部为NaN或不存在，将被排除")
        
        # 添加电池室温度（如果有）
        if chamber_temp is not None and not np.all(np.isnan(chamber_temp)):
            features.append(chamber_temp)
            feature_names.append('Chamber_Temp')
        else:
            print("警告: 电池室温度数据全部为NaN或不存在，将被排除")
        
        # 检查是否有可用特征
        if len(features) == 0:
            raise ValueError("没有可用的特征数据，所有特征都是NaN")
        
        # 转换为特征矩阵
        X = np.column_stack(features)
        y = soc
        
        print(f"Data processed successfully. Features: {X.shape}, SOC: {y.shape}")
        
        return X, y
        
    except Exception as e:
        print(f"Error loading Panasonic data: {e}")
        import traceback
        traceback.print_exc()
        raise


def _extract_from_struct(struct_array, field_name):
    """
    从MATLAB结构体数组中提取字段值
    
    Args:
        struct_array: MATLAB结构体数组
        field_name: 字段名
        
    Returns:
        提取的数组或None（如果未找到）
    """
    try:
        # 检查字段是否存在
        if field_name in struct_array.dtype.names:
            # 提取字段值（通常是一个包含一个元素的数组）
            field_value = struct_array[field_name][0, 0]
            
            # 如果是数组，直接返回
            if isinstance(field_value, np.ndarray):
                # 如果是数值数组，直接展平并返回
                if np.issubdtype(field_value.dtype, np.number):
                    return field_value.flatten()
                
                # 如果是对象数组或字符串数组，返回原始形状
                return field_value
            
            # 如果是cell数组，需要特殊处理
            if field_value.dtype.kind == 'O':
                # 尝试将cell数组转换为数值数组
                try:
                    return np.array([float(x[0]) if x.size > 0 else np.nan for x in field_value.flatten()])
                except:
                    # 如果转换失败，返回原始cell数组
                    return field_value
            
            return field_value
        return None
    except Exception as e:
        print(f"提取字段 {field_name} 时出错: {e}")
        return None


def _extract_array(data: Dict, possible_keys: List[str]) -> np.ndarray:
    """
    Extract array from data dictionary using possible keys
    
    Args:
        data: Dictionary containing data
        possible_keys: List of possible keys to try
        
    Returns:
        Extracted array or None if not found
    """
    for key in possible_keys:
        if key in data and isinstance(data[key], np.ndarray):
            return data[key].flatten()
    return None


def _estimate_soc_from_voltage(voltage: float) -> float:
    """
    Estimate SOC from voltage for Panasonic 18650PF
    Simple linear approximation based on voltage ranges
    
    Args:
        voltage: Battery voltage
        
    Returns:
        Estimated SOC (0-100%)
    """
    # Voltage ranges for Panasonic 18650PF
    if voltage >= 4.15:
        return 100.0
    elif voltage >= 3.9:
        return 90.0 + (voltage - 3.9) * 10.0 / 0.25
    elif voltage >= 3.7:
        return 70.0 + (voltage - 3.7) * 20.0 / 0.2
    elif voltage >= 3.5:
        return 40.0 + (voltage - 3.5) * 30.0 / 0.2
    elif voltage >= 3.3:
        return 20.0 + (voltage - 3.3) * 20.0 / 0.2
    elif voltage >= 3.0:
        return (voltage - 3.0) * 20.0 / 0.3
    else:
        return 0.0


def generate_synthetic_data(n_samples: int = 1000) -> Tuple[np.ndarray, np.ndarray]:
    """
    Generate synthetic battery data for testing
    
    Args:
        n_samples: Number of samples to generate
        
    Returns:
        Tuple of (features, SOC)
    """
    print("Generating synthetic battery data...")
    
    # Time steps
    time = np.linspace(0, 3600, n_samples)  # 1 hour in seconds
    
    # Generate SOC values (decreasing from 100% to 20%)
    soc = 100 - 80 * (time / 3600) + np.random.normal(0, 1, n_samples)
    soc = np.clip(soc, 0, 100)  # Ensure SOC is between 0-100%
    
    # Generate voltage (decreases with SOC)
    # Typical Li-ion battery: 4.2V (full) to 3.0V (empty)
    voltage = 3.0 + 1.2 * (soc / 100) + np.random.normal(0, 0.05, n_samples)
    
    # Generate current (negative during discharge)
    # Add some variations to simulate different loads
    base_current = -2.0  # 2A discharge
    current = base_current + 0.5 * np.sin(time / 180) + np.random.normal(0, 0.2, n_samples)
    
    # Generate temperature (increases during discharge)
    temp = 25 + 10 * (1 - soc / 100) + np.random.normal(0, 1, n_samples)
    
    # Generate internal resistance (increases as battery discharges)
    resistance = 0.1 + 0.05 * (1 - soc / 100) + np.random.normal(0, 0.01, n_samples)
    
    # Calculate power
    power = voltage * current
    
    # Combine features
    X = np.column_stack([voltage, current, temp, resistance, power])
    y = soc
    
    print(f"Synthetic data generated. Features shape: {X.shape}, SOC shape: {y.shape}")
    
    return X, y


def visualize_battery_data(X: np.ndarray, y: np.ndarray, feature_names: List[str]) -> None:
    """
    Visualize battery data
    
    Args:
        X: Feature matrix
        y: Target vector (SOC)
        feature_names: Names of features
    """
    # Create time index
    time_index = np.arange(len(y))
    
    # Plot SOC
    plt.figure(figsize=(12, 6))
    plt.plot(time_index, y)
    plt.title('Battery SOC')
    plt.xlabel('Time Step')
    plt.ylabel('SOC (%)')
    plt.grid(True)
    plt.savefig('battery_soc.png')
    
    # Plot features
    fig, axes = plt.subplots(len(feature_names), 1, figsize=(12, 4 * len(feature_names)))
    
    for i, (feature, name) in enumerate(zip(X.T, feature_names)):
        if len(feature_names) == 1:
            ax = axes
        else:
            ax = axes[i]
        ax.plot(time_index, feature)
        ax.set_title(f'Feature: {name}')
        ax.set_xlabel('Time Step')
        ax.set_ylabel(name)
        ax.grid(True)
    
    plt.tight_layout()
    plt.savefig('battery_features.png')
    
    # Plot correlation between features and SOC
    plt.figure(figsize=(10, 8))
    for i, (feature, name) in enumerate(zip(X.T, feature_names)):
        plt.subplot(2, 3, i + 1)
        plt.scatter(feature, y, alpha=0.5)
        plt.title(f'{name} vs SOC')
        plt.xlabel(name)
        plt.ylabel('SOC (%)')
        plt.grid(True)
    
    plt.tight_layout()
    plt.savefig('feature_correlations.png')
    
    plt.show()